# @Time : 2021/9/28 10:55
# @Author : xiashuobad
# @File : image_utils.py
# @Software: PyCharm
########################
#  _ooOoo_
# o8888888o
# 88" . "88
# (| -_- |)
# O\ = /O
# ___/`---'\____
# .   ' \\| |// `.
# / \\||| : |||// \
# / _||||| -:- |||||- \
# | | \\\ - /// | |
# | \_| ''\---/'' | |
# \ .-\__ `-` ___/-. /
# ___`. .' /--.--\ `. . __
# ."" '< `.___\_<|>_/___.' >'"".
# | | : `- \`.;`\ _ /`;.`/ - ` : | |
# \ \ `-. \_ __\ /__ _/ .-` / /
# ======`-.____`-.___\_____/___.-`____.-'======
# '=---=' bug泛滥 佛已瘫痪
import contextlib
import json
import re
import mmcv
import warnings
import numpy as np
from rasterio.features import sieve
from rasterio.mask import mask
from PIL import Image
from osgeo import ogr, osr
from skimage import morphology
from tqdm import tqdm
from functools import reduce
from rasterio import features

from constant import ImageType, DICT_FIELD_TYPE
from file_utils import MyPath, get_list_suffix_from_dirs
# from tools.parallel import parallel_run
from open_image import OpenImage
# from tools.postprocess.rdp import rdp
from transform_types import reshape_as_raster, adjust_transform, geo_transform2affine, pixel2geo, \
    format_colormap2pil, format_class_info


def read_image_as_array(img_path: str, backend='gdal', **kwargs):
    """
    gdal读取图像，返回图像array以及坐标和投影
    :param img_path: 图片路径
    :param backend：使用的后端图片库，支持cv2，pil，gdal，rasterio
    :return:
    """
    with OpenImage(img_path, backend=backend, **kwargs) as img:
        img_array = img.img_array
    return img_array


def save_mask2color(_mask: (np.ndarray, str), save_path, color_map: (list, dict), dict_replace_values=None,
                    backend='pil', geo_transform=None, project=None):
    """
    将标签图像array或者预测图像array转换成伪彩色图像并保存
    :param _mask: mask图像路径或者array
    :param save_path: 伪彩色标签保存路径
    :param color_map: rgb颜色映射表,支持pil，rasterio以及二维列表三个格式
    :param dict_replace_values: 字典类型！key为原始值(多个值可以用元组)，value为新值(长度为2的元组，第一个元素为key值对应的新值，
    第二个值表示除key值外的其他值对应的新值,默认为None表示不变)'
    :param backend: 使用的后端图片库，'pil','gdal','rasterio'
    :param geo_transform:坐标系数
    :param project:坐标系
    :return:
    """
    if dict_replace_values is None:
        dict_replace_values = {}
    if isinstance(_mask, str):  # 如果mask参数传的是文件路径，则先读取图像
        _mask = read_image_as_array(_mask, backend=backend)
    if len(_mask.shape) == 3:  # 非单通道灰度图，取第一个通道转为灰度图
        _mask = reshape_as_raster(_mask)[0, ...]
    assert len(_mask.shape) == 2, '只支持单通道灰度图！'
    _mask = img_replace_values(_mask, dict_replace_values)
    _mask = _mask.astype(np.uint8)  # 转uint8图片格式
    write_image(_mask, save_path, geo_transform, project, color_map=color_map, backend=backend)


def image_slide_cutting(end_x, end_y, start_x=0, start_y=0, window_size=512, overlapping_pixels=0):
    """
    大影像滑动窗口切割生成器，返回当前切割得图片左上角得像素坐标(x,y)和实际切割的图像的尺寸(window_size_x, window_size_y)
    :param end_x: x上界
    :param end_y: y上界
    :param start_x: x下界
    :param start_y: y下界
    :param window_size: 切割的图片尺寸大小,(宽，高)，为整数时，宽高相同
    :param overlapping_pixels: 窗口滑动时，重叠的像素大小。
    :return:
    """
    assert end_x > start_x and end_y > start_y
    if isinstance(window_size, int):
        window_size_x = window_size_y = window_size
    else:
        assert isinstance(window_size, (list, tuple))
        window_size_x, window_size_y = window_size[:2]
    # 存在切割尺寸大于原图尺寸的情况，特殊处理
    window_size_y = min(window_size_y, end_y - start_y)
    window_size_x = min(window_size_x, end_x - start_x)
    stride_x = window_size_x - overlapping_pixels
    stride_y = window_size_y - overlapping_pixels
    # 按照从左往右，从上往下，逐行扫描
    for y in list(range(start_y, end_y - window_size_y, stride_y)) + [end_y - window_size_y]:
        for x in list(range(start_x, end_x - window_size_x, stride_x)) + [end_x - window_size_x]:
            yield (x, y), (window_size_x, window_size_y)


def slice_generator(image: (str, OpenImage), window_size: (int, tuple, list), overlapping_pixels=0,
                    start_coordinate: (tuple, list) = (0, 0), end_coordinate: (tuple, list) = None) -> (
        np.ndarray, (int, int), (int, int)):
    """
    图像滑动窗口切割生成器，返回当前切割得图片左上角得像素坐标(x,y)和实际切割的图像的尺寸(window_size_x, window_size_y)
    :param image: 图片array
    :param window_size: 切割的图片尺寸大小,(宽，高)，为整数时，宽高相同
    :param overlapping_pixels: 窗口滑动时，重叠的像素大小。
    :param start_coordinate:
    :param end_coordinate:
    :return: 返回切片的array,(x,y),(window_size_x, window_size_y)
    """
    if isinstance(image, str):
        image = OpenImage(image)
    assert isinstance(image, OpenImage), 'image参数应为图片路径或者OpenImage对象！'
    end_coordinate = end_coordinate or (image.width, image.height)

    for (x, y), (window_size_x, window_size_y) in tqdm(
            list(image_slide_cutting(*end_coordinate, *start_coordinate, window_size,
                                     overlapping_pixels))):
        yield image.read_window(x, y, window_size_x, window_size_y), (x, y), (window_size_x, window_size_y)
    image.close()


def create_field_defn(out_layer, class_fields: dict):
    """
    给据所给字段属性列表给图层创建属性字段
    Args:
        out_layer: 图层
        class_fields: 类别字段，{字段名称:字段类型}

    Returns:

    """
    for field, field_type in class_fields.items():
        fd = ogr.FieldDefn(field, DICT_FIELD_TYPE.get(field_type, ogr.OFTString))
        out_layer.CreateField(fd)


def shp2raster(raster_path: str, shp: (str, list), save_tif_path, field_raster: str, field_mapping=None,
               color_mapping=None, attr_filter: str = None):
    """
    shp 转 栅格图像
    :param raster_path: tif 原图，用来提供坐标信息
    :param shp: 待转的 shp文件路径或者postgis数据库连接地址,或者直接shapes对象:[(geometry,id),...]
    :param save_tif_path: 转好的 tif 保存路径
    :param field_raster: 用来匹配的字段名称
    :param field_mapping: 用来匹配信息的字段映射，dict
    :param color_mapping: rgb颜色映射表，rasterio格式或pil格式
    :param attr_filter: 属性过滤条件查询语句
    """
    if field_mapping is None:
        field_mapping = {}
    with OpenImage(raster_path) as img_ds:
        geo_transform = img_ds.geo_transform
        projection = img_ds.project
        width, height = img_ds.size
    geo_transform = geo_transform2affine(geo_transform)  # 转换为rasterio格式
    if isinstance(shp, str):
        with read_shp(shp, attr_filter) as layer:
            shp = [(json.loads(feature.geometry().ExportToJson()),
                    field_mapping[feature.GetField(field_raster)] if field_mapping else feature.GetField(field_raster))
                   for feature in layer]

    mask_array = features.rasterize(shp, (height, width), transform=geo_transform, all_touched=True)
    write_image(mask_array, save_tif_path, geo_transform, projection, color_map=color_mapping)


@contextlib.contextmanager
def read_shp(shp_path: str, attr_filter=None, spatial_filter=None) -> ogr.Layer:
    """
    读取shp，返回layer，支持shpfile文件路径和postgis数据库读取
    :param shp_path: 数据库连接地址格式为 "PG:host=localhost port=5432 dbname=kqgis user=postgres password=135456
    table=dawenhe-clip name='水域' or name='房屋建筑'"
    数据库连接地址格式分为3个部分："PG:host=localhost port=5432 dbname=kqgis user=postgres password=135456","table=...",
    "name=...",后面两个分别表示查询的表名和属性过滤条件
    :param attr_filter: 属性过滤条件表达式，字符串，例如："name='水域' or name='房屋建筑'"
    :param spatial_filter: 空间过滤，Geometry类型，给定空间范围
    :return:
    """
    assert isinstance(shp_path, str)
    if MyPath(shp_path).is_file():  # shp文件路径
        ds: ogr.DataSource = ogr.Open(shp_path)
        layer = ds.GetLayer()
    elif shp_path.startswith('PG:') and 'table' in shp_path:  # postgis数据库地址
        res = re.findall(r'\s*(PG:.*)\s+table=(\S+)\s*', shp_path.strip())
        assert len(res) > 0, f'{shp_path}格式错误！'
        db_path, table = res[0]
        ds: ogr.DataSource = ogr.Open(db_path)
        layer: ogr.Layer = ds.GetLayerByName(table)
    else:
        raise Exception(f'{shp_path} 格式错误！')
    attr_filter and layer.SetAttributeFilter(attr_filter)
    spatial_filter and layer.SetSpatialFilter(spatial_filter)
    assert len(layer) > 0, f'{shp_path} 读取错误！'
    yield layer
    del ds


def write_image(img: (np.ndarray, str), save_path, geo_transform=None, projection_ref=None, ref_raster_path=None,
                color_map: dict = None, nodata=None, backend='gdal'):
    """
    图像array保存为带坐标的tif栅格数据
    :param img:
    :param save_path: 保存路径
    :param geo_transform: 地理坐标变换系数
    :param projection_ref: 坐标系
    :param ref_raster_path: 参考栅格影像路径
    :param color_map: rgb颜色映射表，字典类型
    :param nodata: 无效值
    :param backend: 'gdal','rasterio','pil','cv2'
    :return:
    """
    if isinstance(img, str):  # 传入的是图片路径时，先读取图片
        with OpenImage(img, backend=backend) as ds_img:
            img = ds_img.img_array
            geo_transform = geo_transform or ds_img.geo_transform  # 当参数设置了geo_transform的值就以参数，否则以读取图片里的值
            projection_ref = projection_ref or ds_img.project  # 同上
            nodata = nodata or ds_img.nodata  # 同上
    if ref_raster_path and MyPath(ref_raster_path).exists():  # 当设置了参考栅格的路径时优先级最高，使用参考栅格的坐标变换系数和坐标系
        with OpenImage(ref_raster_path, backend=backend) as ds_img:
            geo_transform = ds_img.geo_transform
            projection_ref = ds_img.project
            nodata = ds_img.nodata
    MyPath(save_path).make_parent_dirs()  # 提前创建父目录，防止写入文件报错
    with OpenImage(save_path, 'w', backend) as ds_img:  # 写入数据到图片文件
        ds_img.set_img_array(img).set_nodata(nodata).set_geo_transform(geo_transform).set_project(
            projection_ref).set_color_map(color_map)


def mask2shp(_mask: (np.ndarray, str), save_shp_path: str, class_info: (list, dict), geo_transform=None,
             projection_ref=None, ref_raster_path=None, del_background=False, simplify_tolerance=2):
    """
    mask预测图array转矢量shp
    :param _mask: mask图片路径或者array
    :param save_shp_path: 矢量shp文件保存路径
    :param class_info: 分类信息，格式同dict_model配置文件
    :param geo_transform: 地理坐标变换系数
    :param projection_ref: 投影信息
    :param ref_raster_path: 参考影像
    :param del_background: 是否删除背景0
    :param simplify_tolerance: 边缘点精简，容忍度参数，值越大精简越狠
    :return: 返回所有feature（目标区域）的列表
    """
    if not isinstance(class_info, dict):
        class_info = format_class_info(class_info)
    if isinstance(_mask, str):  # mask为文件路径时，先读取mask图片文件
        with OpenImage(_mask) as img:
            _mask = img.img_array
            geo_transform = geo_transform or img.geo_transform  # 参数未设置地理坐标仿射变换和坐标系时，直接从mask图像文件中获取
            projection_ref = projection_ref or img.project
    assert len(_mask.shape) == 2, 'mask_array应为单通道灰度图！'
    if ref_raster_path and MyPath(ref_raster_path).exists():  # 当设置了参考栅格影像的路径的时，地理坐标和坐标系以参考栅格影像的为准
        with OpenImage(ref_raster_path) as img:
            geo_transform = img.geo_transform
            projection_ref = img.project
    MyPath(save_shp_path).make_parent_dirs()  # 自动创建结果矢量的目录
    geo_transform = adjust_transform(geo_transform)  # 无坐标系默认变换系数的情况下会上下翻转，自动调整
    geo_transform = geo_transform2affine(geo_transform)  # 转换为rasterio支持的格式affine
    shapes = features.shapes(_mask, mask=(_mask != 0) if del_background else None,
                             transform=geo_transform)  # 栅格矢量化
    drv: ogr.Driver = ogr.GetDriverByName('ESRI Shapefile')  # 创建矢量驱动
    if MyPath(save_shp_path).exists():  # 删除之前的结果
        drv.DeleteDataSource(save_shp_path)  # 这个函数会将.shp文件以及它的依赖文件一起删除
    dst_shp_ds = drv.CreateDataSource(save_shp_path)  # 创建文件资源
    proj = osr.SpatialReference()  # 创建坐标系类
    proj.ImportFromWkt(projection_ref if projection_ref else '')  # 导入坐标系
    dst_layer: ogr.Layer = dst_shp_ds.CreateLayer('layer', geom_type=ogr.wkbPolygon, srs=proj)  # 创建多边形图层，并写入坐标系
    create_field_defn(dst_layer, class_fields := class_info['class_fields'])  # 创建矢量属性字段
    for polygon, class_id in shapes:  # 遍历shapes，每个shape中保存的是要素的类型以及坐标（字典格式），还有id值
        class_id = int(class_id)
        geometry: ogr.Geometry = ogr.CreateGeometryFromJson(json.dumps(polygon))  # 先将字典转为json格式，再从geojson格式创建geometry
        # 简化边缘点
        geometry = geometry.SimplifyPreserveTopology(simplify_tolerance * geo_transform.a)
        feature: ogr.Feature = ogr.Feature(dst_layer.GetLayerDefn())  # 创建feature要素
        for field, field_type in class_fields.items():  # 遍历所有属性字段，逐个设置值
            # 当类别列表中未设置该属性字段的值时，根据矢量字段类型设置默认值
            field_value = class_info['class_names'][class_id].get(field, field_type())
            feature.SetField(field, field_value)  # 设置字段的值
        feature.SetGeometry(geometry)  # 给feature要素添加多边形的坐标点位置信息
        dst_layer.CreateFeature(feature)  # 将feature要素加到图层中

    dst_shp_ds.Destroy()  # 释放资源，更新写入shp文件


def stitching_image(list_img_paths, save_path=None):
    """
    将多张图片水平拼接成一张图片
    :param list_img_paths: 待拼接的图片路劲列表
    :param save_path: 拼接后的图片保存路径
    :return: 返回拼接后的图片array
    """
    list_img_arrays = [read_image_as_array(img_path, backend='cv2') for img_path in
                       list_img_paths]  # 顺序读取所有图片，array存入列表
    img_stitched = np.hstack(list_img_arrays)  # 将所有图片进行水平拼接
    save_path and write_image(img_stitched, save_path, backend='cv2')  # 保存路径不为空，则保存到文件
    return img_stitched


def write_coordinate_to_raster(target_image_path, geo_transform=None, project=None, reference_image_path=None):
    """
    给目标影像写入地理坐标
    :param reference_image_path:参考图片路径
    :param target_image_path:目标图片路径
    :param geo_transform:坐标变换系数
    :param project:坐标系
    :return:
    """
    if reference_image_path is not None:
        with OpenImage(reference_image_path) as reference_data:
            geo_transform = reference_data.geo_transform
            project = reference_data.project
    with OpenImage(target_image_path, 'u') as target_data:
        target_data.set_geo_transform(geo_transform).set_project(project)


def show_image(image: (np.ndarray, str), color_map=None):
    """
    简单显示一下图片，单通道灰度图可以设置rgb映射
    Args:
        image: 图片路径或者array
        color_map: rgb映射列表

    Returns:

    """
    if isinstance(image, str):
        image = read_image_as_array(image)
    image = Image.fromarray(image)
    if color_map is not None:
        color_map = format_colormap2pil(color_map)
        image.putpalette(color_map)
    image.show()


def generate_graph(im):
    if len(im.shape) == 3:
        warnings.warn(f'图片尺寸为：{im.shape}，不是单通道，将只取第一个通道。')
        im = im[:, :, 0] if np.array(im.shape).argmin() == 2 else im[0, :, :]
    im = np.swapaxes(im, 0, 1)
    # im_temp = im.copy()
    im[im > 0] = 1

    im = morphology.skeletonize(im)  # * im_temp
    im = im.astype(np.uint8)

    # extract a graph by placing vertices every THRESHOLD pixels, and at all intersections
    vertices = []
    edges = set()

    def add_edge(src, dst):
        if (src, dst) in edges or (dst, src) in edges:
            return
        elif src == dst:
            return
        edges.add((src, dst))

    point_to_neighbors = {}
    q = []
    while True:
        if len(q) > 0:
            last_id, i, j = q.pop()
            path = [vertices[last_id], (i, j)]
            if im[i, j] == 0:
                continue
            point_to_neighbors[(i, j)].remove(last_id)
            if len(point_to_neighbors[(i, j)]) == 0:
                del point_to_neighbors[(i, j)]
        else:
            w = np.where(im > 0)
            if len(w[0]) == 0:
                break
            i, j = w[0][0], w[1][0]
            last_id = len(vertices)
            vertices.append((i, j))
            path = [(i, j)]

        while True:
            im[i, j] = 0
            neighbors = []
            for oi in [-1, 0, 1]:
                for oj in [-1, 0, 1]:
                    ni = i + oi
                    nj = j + oj
                    if 0 <= ni < im.shape[0] and 0 <= nj < im.shape[1] and im[ni, nj] > 0:
                        neighbors.append((ni, nj))
            if len(neighbors) == 1 and (i, j) not in point_to_neighbors:
                ni, nj = neighbors[0]
                path.append((ni, nj))
                i, j = ni, nj
            else:
                if len(path) > 1:
                    # path = rdp(path, 2)
                    if len(path) > 2:
                        for point in path[1:-1]:
                            cur_id = len(vertices)
                            vertices.append(point)
                            add_edge(last_id, cur_id)
                            last_id = cur_id
                    neighbor_count = len(neighbors) + len(point_to_neighbors.get((i, j), []))
                    if neighbor_count == 0 or neighbor_count >= 2:
                        cur_id = len(vertices)
                        vertices.append(path[-1])
                        add_edge(last_id, cur_id)
                        last_id = cur_id
                for ni, nj in neighbors:
                    if (ni, nj) not in point_to_neighbors:
                        point_to_neighbors[(ni, nj)] = set()
                    point_to_neighbors[(ni, nj)].add(last_id)
                    q.append((last_id, ni, nj))
                for neighbor_id in point_to_neighbors.get((i, j), []):
                    add_edge(neighbor_id, last_id)
                break
    neighbors = {}

    vertex = vertices

    for edge in edges:

        nk1 = (vertex[edge[0]][1], vertex[edge[0]][0])
        nk2 = (vertex[edge[1]][1], vertex[edge[1]][0])

        if nk1 != nk2:
            if nk1 in neighbors:
                if nk2 in neighbors[nk1]:
                    pass
                else:
                    neighbors[nk1].append(nk2)
            else:
                neighbors[nk1] = [nk2]

            if nk2 in neighbors:
                if nk1 in neighbors[nk2]:
                    pass
                else:
                    neighbors[nk2].append(nk1)
            else:
                neighbors[nk2] = [nk1]

    return neighbors


def raster2line_shp(_mask: np.ndarray, geo_trans, project_ref, out_line_shp_path, class_infos):
    if not isinstance(class_infos, dict):
        class_infos = format_class_info(class_infos)
    assert len(_mask.shape) == 2, '必须是单通道影像！'
    shp_driver = ogr.GetDriverByName("ESRI Shapefile")
    MyPath(out_line_shp_path).make_parent_dirs()
    if MyPath(out_line_shp_path).exists():
        shp_driver.DeleteDataSource(out_line_shp_path)
    out_datasource = shp_driver.CreateDataSource(out_line_shp_path)
    geo_trans = adjust_transform(geo_trans)
    proj = osr.SpatialReference()
    proj.ImportFromWkt(project_ref)
    out_layer: ogr.Layer = out_datasource.CreateLayer('layer', srs=proj, geom_type=ogr.wkbMultiLineString)
    create_field_defn(out_layer, class_fields := class_infos['class_fields'])
    out_defn = out_layer.GetLayerDefn()  # 定义要素
    out_feature_triangle = ogr.Feature(out_defn)
    for i, class_info in enumerate(class_infos):
        if i == 0:  # 背景
            continue
        cur_mask = (_mask == i).astype(np.uint8)
        graph = generate_graph(cur_mask)
        # 创建单个面
        for n, v in graph.items():
            for nei in v:
                line = ogr.Geometry(ogr.wkbLinearRing)  # 构建几何类型:线
                nx, ny = n[1], n[0]
                nx, ny = pixel2geo(nx, ny, geo_trans)
                line.AddPoint(nx, ny)  # 添加点01

                nei_x, nei_y = nei[1], nei[0]
                nei_x, nei_y = pixel2geo(nei_x, nei_y, geo_trans)
                line.AddPoint(nei_x, nei_y)  # 添加点02

                out_feature_triangle.SetGeometry(line)
                for field, field_type in class_fields.items():
                    out_feature_triangle.SetField(field, class_info['class_names'].get(field, field_type(0)))
                out_layer.CreateFeature(out_feature_triangle)

    out_datasource.Destroy()


def higher_bit2img8bit(img_higher_bit: (str, np.ndarray), img8bit_path=None, to_rgb=False, geo_transform=None,
                       proj=None, nodata=None):
    """
    遥感影像高位（16位，32位。。。）转8位影像
    Args:
        img_higher_bit: 待转换的图像
        img8bit_path: 转换后的8位图像保存地址
        to_rgb:
        geo_transform:
        proj:
        nodata:
    Returns:

    """
    if isinstance(img_higher_bit, str):  # 如果传进来的参数是图片路径，则先读取图片
        with OpenImage(img_higher_bit) as ds:
            nodata = ds.nodata  # 获取无效值
            geo_transform, proj = ds.geo_transform, ds.project  # 地理坐标变换系数和坐标系
            array_data = ds.img_array  # 读取array
    else:
        assert isinstance(img_higher_bit, np.ndarray)
        array_data = img_higher_bit
    if array_data.dtype == 'uint8':  # 如果本身就是8位，直接返回
        return array_data
    array_data[array_data == nodata] = 0  # 将无效值置为0
    if len(array_data.shape) == 2:  # 单通道灰度图，先升维，以便和彩色影像同样处理
        array_data = array_data[None, ...]
    counts, rows, cols = array_data.shape
    if to_rgb and counts > 3:  # to_rgb为真时，要求影像波段数大于3，只取前三个波段，RGB
        counts = 3
        array_data = array_data[:3, ...][::-1, ...]
    compress_data = np.zeros((counts, rows, cols), dtype=np.uint8)  # 初始化一个最终的8位影像

    for i in range(counts):  # 遍历各个波段，分别进行灰度值的压缩

        cur_band = array_data[i, ...]
        cut_min, cut_max = np.percentile(cur_band, 2), np.percentile(cur_band, 98)
        if cut_min == cut_max:
            continue
        compress_scale = (cut_max - cut_min) / 255  # 灰度值缩放比例
        cur_band[cur_band > cut_max] = cut_max  # 大于上区间的值直接过滤掉，下同
        cur_band[cur_band < cut_min] = cut_min
        compress_data[i, :, :] = (cur_band - cut_min) / compress_scale  # 将灰度值缩放到8位图的[0，255]之间

    if counts == 1:  # 单波段影像还原
        compress_data = compress_data[0, ...]

    if img8bit_path is not None:  # 设置了8位图保存路径则保存到文件，否则不保存
        write_image(compress_data, img8bit_path, geo_transform, proj)  # 将转换后的8位图写入文件
    return compress_data


def resize_img(img: (str, Image.Image), target_size: tuple, save_path=None, pad=None):
    """
    使用Image对图片进行resize
    Args:
        img:
        target_size:(w,h)
        save_path:
        pad:
    Returns:

    """
    if isinstance(img, str):
        img: Image.Image = Image.open(img)
    assert isinstance(img, Image.Image)
    ori_w, ori_h = ori_size = img.size
    new_w, new_h = target_size
    if ori_size == target_size:
        return img
    if pad is None:  # 不填充，直接拉伸
        new_img = img.resize(target_size, Image.NEAREST)
    else:  # 保持横纵比，其余部分填充
        assert isinstance(pad, int), 'pad参数应该为整型数值，表示要填充的像素值'
        new_img = Image.new(img.mode, target_size, pad)
        scale_rate_w, scale_rate_h = new_w / ori_w, new_h / ori_h
        if scale_rate_w < scale_rate_h:
            img = img.resize((new_w, cur_h := int(ori_h * scale_rate_w)))
            new_img.paste(img, (0, (new_h - cur_h) // 2))
        else:
            img = img.resize((cur_w := int(ori_w * scale_rate_h), new_h))
            new_img.paste(img, ((new_w - cur_w) // 2, 0))
    if save_path is not None:
        MyPath(save_path).make_parent_dirs()
        new_img.save(save_path)
        print(f'{save_path}尺寸已从{ori_size}resize为{target_size}')
    return new_img


def resize_img_group(images: (list, tuple), target_size: tuple, sava_paths=None, pad=None):
    """
    使用Image对一组图片进行同步resize
    Args:
        images:
        target_size:
        sava_paths:
        pad:
    Returns:

    """
    if isinstance(images, str):  # 单张图片
        images = [images]
    assert isinstance(images, (list, tuple))
    results = []
    if sava_paths is not None:
        if isinstance(sava_paths, str):
            sava_paths = [sava_paths]
        assert isinstance(sava_paths, (list, tuple)) and len(images) == len(sava_paths)
    else:
        sava_paths = [None] * len(images)

    for img, save_path in zip(images, sava_paths):
        img = resize_img(img, target_size, save_path, pad)
        results.append(img)

    return results


# gdal版本
# def extract_by_shp_gdal(in_shp_path, in_raster_path, out_raster_path='mem', nodata=0):
#     """根据矢量的不规则范围对栅格影像进行裁剪，非范围内的赋无效值"""
#     input_raster = gdal.Open(in_raster_path)
#
#     # 利用gdal.Warp进行裁剪
#     # https://gdal.org/api/python/osgeo.gdal.html#osgeo.gdal.Warp
#     MyPath(out_raster_path).make_parent_dirs()
#     result: gdal.Dataset = gdal.Warp(
#         out_raster_path,
#         input_raster,
#         format='GTiff',
#         cutlineDSName=in_shp_path,  # 用于裁剪的矢量
#         cropToCutline=True,  # 是否使用cut_lineDSName的extent作为输出的界线
#         dstNodata=nodata  # 输出数据的nodata值
#     )
#     result.FlushCache()
#     out_image_array = result.ReadAsArray()
#     out_transform = result.GetGeoTransform()
#     del result
#     return out_image_array, out_transform


# rasterio版本
def extract_by_shp(in_shp, in_raster_path, out_raster_path=None):
    """
    根据矢量的不规则范围对栅格影像进行裁剪，非范围内的赋无效值
    :param in_shp: shp文件路径或者wkt格式的polygon
    :param in_raster_path: 图片路径
    :param out_raster_path: 剪裁图片保存路径
    :return:
    """
    assert isinstance(in_shp, str), (error_in_shp := 'inshp参数只能是shp文件路径或者wkt格式的polygon字符串！')
    with OpenImage(in_raster_path, backend='rasterio') as src:
        if MyPath(in_shp).is_file():
            # x_min, y_min, x_max, y_max = src.img_data.bounds
            # wkt = f"POLYGON (({x_min} {y_min},{x_max} {y_min},{x_max} {y_max},{x_min} {y_max}))"
            # spatial_filter = ogr.CreateGeometryFromWkt(wkt)
            with read_shp(in_shp, spatial_filter=None) as layer:
                in_shp = [json.loads(feature.geometry().ExportToJson()) for feature in layer]
        elif in_shp.startswith('POLYGON'):
            geometry_bound: ogr.Geometry = ogr.CreateGeometryFromWkt(in_shp)
            in_shp = [json.loads(geometry_bound.ExportToJson())]
        else:
            raise ValueError(error_in_shp)
        nodata = src.nodata or 0
        out_image, out_transform = mask(src.img_data, in_shp, nodata=nodata, crop=True)
        geo_project = src.project

    # 保存
    if out_raster_path is not None:
        MyPath(out_raster_path).make_parent_dirs()
        write_image(out_image, out_raster_path, out_transform, geo_project, nodata=nodata)
    return out_image, out_transform


def rgb_label2mask(rgb_label_path: str, rgb_mapping: dict, out_path=None, palette=None):
    """
    rgb真彩色标签转单通道训练标签
    :param rgb_label_path: rgb标签图片路径，
    :param rgb_mapping: rgb颜色映射字典，{(r,g,b):class_id}，key,value分别对应rgb值和类别id值
    :param out_path: 标签保存路径，默认为空
    :param palette: 伪彩rgb映射列表
    :return:
    """
    rgb_label = read_image_as_array(rgb_label_path, backend='pil')
    assert len(rgb_label.shape) == 3, '必须是真彩色标签图像！'
    out_label = np.zeros(rgb_label.shape[:2], dtype=np.uint8)
    for rgb, class_id in rgb_mapping.items():
        r, g, b = rgb[:3]
        mask_rgb = (rgb_label[..., 0] == r) & (rgb_label[..., 1] == g) & (rgb_label[..., 2] == b)
        out_label[mask_rgb] = class_id
    if out_path is None:
        return out_label
    MyPath(out_path).make_parent_dirs()
    with OpenImage(out_path, 'w', backend='pil') as out:
        out.set_img_array(out_label)
        if palette is None:
            palette = {class_id: rgb for rgb, class_id in rgb_mapping.items()}
            palette = np.array(list(dict(sorted(palette.items())).values())).flatten().tolist()
        out.set_color_map(palette)
    return out_label


def batch_rgb_label2mask(rgb_label_dir, rgb_mapping: dict, out_dir, palette=None):
    for label_path in tqdm(MyPath(rgb_label_dir).iter_dir()):
        out_path = label_path.replace_parent(out_dir)
        rgb_label2mask(str(label_path), rgb_mapping, str(out_path), palette)


def seg_result2change_result(seg_result1, seg_result2, dict_replace_values=None):
    if dict_replace_values is None:
        dict_replace_values = {255: 2}
    if isinstance(seg_result1, str):
        seg_result1 = read_image_as_array(seg_result1)
    if isinstance(seg_result2, str):
        seg_result2 = read_image_as_array(seg_result2)
    assert isinstance(seg_result1, np.ndarray) and isinstance(seg_result2, np.ndarray)
    if seg_result1.shape != seg_result2.shape:
        seg_result2 = mmcv.imresize(seg_result2, seg_result1.shape[::-1])

    change_result = seg_result2 - seg_result1
    change_result = sieve(change_result, size=100)
    change_result = img_replace_values(change_result, dict_replace_values)
    return change_result


def img_replace_values(img: (np.ndarray, str), dict_replace_values: dict, save_path=None) -> np.ndarray:
    """
    替换img_array中的值
    Args:
        img: 图片路径或者array
        dict_replace_values: 替换字典，格式例子：{255:0},{2:(1,0)},{(2,3):(1,0)},{(255,1):0,(2,3):1,(4,5):2}
        save_path: 保存路径，默认为None，不保存。

    Returns:

    """

    assert isinstance(dict_replace_values,
                      dict), 'replace_values参数类型应该为字典类型！key为原始值(多个值可以用元组)，value为新值(' \
                             '长度为2的元组，第一个元素为key值对应的新值，第二个值表示除key值外的其他值对应的新值,默认为None表示不变)'
    if isinstance(img, str):
        img = read_image_as_array(img)
    assert isinstance(img, np.ndarray), f'img参数类型错误！'
    for ori_values, new_values in dict_replace_values.items():
        if isinstance(ori_values, int):  # key为单个值整数时，先变为元组
            ori_values = (ori_values,)
        assert isinstance(ori_values, tuple)
        if isinstance(new_values, int):
            new_values = (new_values, None)
        assert isinstance(new_values, (tuple, list)) and len(new_values) >= 2
        if len(ori_values) == 1:  # 要改的值只有一个时，不适用于reduce函数
            arg_mask = img == ori_values[0]
        else:  # 同时修改多个值，用reduce函数生成掩码
            arg_mask = reduce(lambda x, y: (img == x) | (img == y), ori_values)
        img = np.where(arg_mask, new_values[0],
                       new_values[1] if new_values[1] is not None else img)  # 字典的value元组第二值为None时，表示其他值不变
    if save_path is not None:
        write_image(img.astype(np.uint8), save_path)
    return img


def mask_valid(img_array):
    """
    计算有效值掩码,即多个通道不全为0或者不全为255的为有效
    :param img_array:
    :return:
    """
    mask_nodata1 = img_array != 0
    mask_nodata2 = img_array != 255
    if len(img_array.shape) == 3:  # 如果是rgb影像，需要对多个通道进行判断有效值，并转为单通道掩码，与最终预测结果shape保持一致
        mask_nodata1 = reduce(lambda _x, _y: _x | _y, mask_nodata1)  # 多个通道取并，即只要任意一个通道不是无效值，就判断为有效
        mask_nodata2 = reduce(lambda _x, _y: _x | _y, mask_nodata2)
    return mask_nodata1 * mask_nodata2
